from pyparsing import (
    Word, alphas, alphanums, Forward, Group, 
    Suppress, Optional as PyparsingOptional, delimitedList, CaselessKeyword,
    ParseException, Literal, ZeroOrMore
)

try:
    from .type_schema import TypeSchema
    from .error import ProtocolError
except ImportError:
    from type_schema import TypeSchema
    from error import ProtocolError

# Define basic type and identifier
VOID = Literal("Void") | Literal("void")
INT = Literal("Int") | Literal("int")
FLOAT = Literal("Float") | Literal("float")
STRING = Literal("String") | Literal("str")
BOOL = Literal("Bool") | Literal("bool")
JSON = Literal("Json") | Literal("json")
identifier = Word(alphas, alphanums + "_")

def create_type_parser():
    """Create type parser"""
    # Create Forward element for recursive reference
    type_expr = Forward()
    
    # Basic types
    basic_type = (
        VOID.copy().setParseAction(lambda: TypeSchema.Void()) |
        INT.copy().setParseAction(lambda: TypeSchema.Int()) |
        FLOAT.copy().setParseAction(lambda: TypeSchema.Float()) |
        STRING.copy().setParseAction(lambda: TypeSchema.Str()) |
        BOOL.copy().setParseAction(lambda: TypeSchema.Bool()) |
        JSON.copy().setParseAction(lambda: TypeSchema.Json())
    )
    
    # List type: List<Type>
    list_type_angle = (
        Suppress("List<") + type_expr + Suppress(">")
    ).setParseAction(lambda tokens: TypeSchema.List(tokens[0]))
    
    # List type: [Type] - optional bracket notation
    list_type_bracket = (
        Suppress("[") + type_expr + Suppress("]")
    ).setParseAction(lambda tokens: TypeSchema.List(tokens[0]))
    
    # Map type: Map<KeyType, ValueType>
    map_type = (
        Suppress("Map<") + type_expr + Suppress(",") + type_expr + Suppress(">")
    ).setParseAction(lambda tokens: TypeSchema.Map(tokens[0], tokens[1]))
    
    # Tuple type: (Type1, Type2, ...) - more robust implementation
    tuple_elements = delimitedList(type_expr)
    tuple_type = (
        Suppress("(") + PyparsingOptional(tuple_elements) + Suppress(")")
    ).setParseAction(lambda tokens: TypeSchema.Tuple(tokens.asList() if tokens else []))

    tuple_type_alt = (
        Suppress("Tuple<") + PyparsingOptional(tuple_elements) + Suppress(">")
    ).setParseAction(lambda tokens: TypeSchema.Tuple(tokens.asList() if tokens else []))
    
    # Optional type: ?Type
    option_type = (
        Suppress("?") + type_expr
    ).setParseAction(lambda tokens: TypeSchema.Opt(tokens[0]))

    option_type_alt = (
        Suppress("Option<") + type_expr + Suppress(">")
    ).setParseAction(lambda tokens: TypeSchema.Opt(tokens[0]))
    
    # Custom type
    custom_type = (
        ~(VOID | INT | FLOAT | STRING | BOOL | JSON | 
          Literal("List") | Literal("Map") | Literal("?")) + 
        identifier.copy()
    ).setParseAction(lambda tokens: TypeSchema.DataType(tokens[0]))
    
    type_expr << (
        basic_type |
        list_type_angle |
        list_type_bracket |
        map_type |
        tuple_type |
        tuple_type_alt |
        option_type |
        option_type_alt |
        custom_type
    )
    
    return type_expr

_type_parser = create_type_parser()


def parse_abs_type(input_str: str) -> TypeSchema:
    """Parse type string into TypeSchema object"""
    if not input_str or not input_str.strip():
        raise ProtocolError(f"Type string cannot be empty")
    clean_input = input_str.replace(" ", "")
    
    try:
        result = _type_parser.parseString(clean_input, parseAll=True)
        return result[0]
    except ParseException as e:
        position = e.loc
        line_info = f"Position {position}:" if position >= 0 else ""
        arrow = " " * position + "^" if position >= 0 else ""
        inner_errmsg = str(e)
        if len(inner_errmsg) > 100:
            inner_errmsg = inner_errmsg[:100] + "..."
        error_msg = f"Type parsing error '{input_str}': {line_info}\n{clean_input}\n{arrow}\n{inner_errmsg}"
        raise ProtocolError(error_msg)

def task_nominal_to_symbolic(task: str) -> tuple[str, TypeSchema, TypeSchema]:
    """Convert nominal task to symbolic task"""
    if not task or not task.strip():
        raise ProtocolError(f"Task signature cannot be empty")
    
    splits = task.split(':')
    if len(splits) != 2:
        raise ProtocolError(f"Task signature cannot be correctly split by colon: {task}")
    
    name, type_sig = splits
    name = name.strip()
    if not name:
        raise ProtocolError(f"Task name cannot be empty: {task}")
    
    type_splits = type_sig.split('->')
    if len(type_splits) != 2:
        raise ProtocolError(f"Task signature cannot be correctly split by arrow: {task}")
    
    input_type_str, output_type_str = type_splits
    input_type_str = input_type_str.strip()
    output_type_str = output_type_str.strip()
    
    try:
        input_type = parse_abs_type(input_type_str)
    except ProtocolError as e:
        raise ProtocolError(f"Task input type parsing error: {task}\nReason: {e}")
    
    try:
        output_type = parse_abs_type(output_type_str)
    except ProtocolError as e:
        raise ProtocolError(f"Task output type parsing error: {task}\nReason: {e}")
    
    return name, input_type, output_type

if __name__ == "__main__":
    print("===== Testing Type Parsing =====")
    test_types = [
        # success
        "(Int, String, Float)",
        "List<Map<String, Int>>",
        "[Map<String, [Int]>]",
        "Map<String, List<(Int, Float)>>",
        "CustomType",
        "Map<(Int, String), [Float]>",
        # fail
        "",
        "List<>",
        "Map<Int>",
        "Map<Int,>",
        "(Int,",
        "[Int",
        "?",
    ]
    for t in test_types:
        try:
            result = parse_abs_type(t)
            print(f"✓ Successfully parsed '{t}' as {result}")
        except ProtocolError as e:
            print(f"✗ Parsing failed for '{t}': {e}")
    print("\n===== Testing Task Parsing =====")
    test_tasks = [
        # success
        "main:FilePath->Map<String,List<List<(Int,Float)>>>",
        "process:Void->String",
        "calculate:Int->Float",
        "analyze:[Int]->(Float, String)",
        # fail
        "",
        "nocolon",
        "name:",
        "name:Int",
        "name:->Int",
        "name:Int->",
        "name:Unknown->Int",
        "name:Int->List<>",
    ]
    for t in test_tasks:
        try:
            task = task_nominal_to_symbolic(t)
            print(f"✓ Successfully parsed task '{t}'")
            print(f"  Name: {task.name}")
            print(f"  Input type: {task.input_type}")
            print(f"  Output type: {task.output_type}")
        except ProtocolError as e:
            print(f"✗ Task parsing failed for '{t}': {e}") 